揭秘 WebGL 计算着色器工作组：深入了解并行处理组织

            #version 300 es
precision highp float;

// Input arrays (global memory)
in  layout(binding = 0) readonly buffer InputA { float inputArrayA[]; };
in  layout(binding = 1) readonly buffer InputB { float inputArrayB[]; };

// Output array (global memory)
out layout(binding = 2) buffer OutputC { float outputArrayC[]; };

// Number of elements per workgroup
layout(local_size_x = 64) in;

// The workgroup ID and local ID are automatically available to the shader.
void main() {
  // Calculate the index within the arrays
  uint index = gl_GlobalInvocationID.x; // Use gl_GlobalInvocationID for global index

  // Add the corresponding elements
  outputArrayC[index] = inputArrayA[index] + inputArrayB[index];
}
            

              
                Copy
              
            
          

            // Get the WebGL context
const canvas = document.createElement('canvas');
document.body.appendChild(canvas);
const gl = canvas.getContext('webgl2');
if (!gl) {
    console.error('WebGL2 not supported');
}

// Shader source
const shaderSource = `#version 300 es
precision highp float;

// Input arrays (global memory)
in  layout(binding = 0) readonly buffer InputA { float inputArrayA[]; };
in  layout(binding = 1) readonly buffer InputB { float inputArrayB[]; };

// Output array (global memory)
out layout(binding = 2) buffer OutputC { float outputArrayC[]; };

// Number of elements per workgroup
layout(local_size_x = 64) in;

// The workgroup ID and local ID are automatically available to the shader.
void main() {
  // Calculate the index within the arrays
  uint index = gl_GlobalInvocationID.x; // Use gl_GlobalInvocationID for global index

  // Add the corresponding elements
  outputArrayC[index] = inputArrayA[index] + inputArrayB[index];
}
`;

// Compile shader
function createShader(gl, type, source) {
    const shader = gl.createShader(type);
    gl.shaderSource(shader, source);
    gl.compileShader(shader);
    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
        console.error('An error occurred compiling the shaders: ' + gl.getShaderInfoLog(shader));
        gl.deleteShader(shader);
        return null;
    }
    return shader;
}

// Create and link the compute program
function createComputeProgram(gl, shaderSource) {
    const computeShader = createShader(gl, gl.COMPUTE_SHADER, shaderSource);
    if (!computeShader) {
        return null;
    }

    const program = gl.createProgram();
    gl.attachShader(program, computeShader);
    gl.linkProgram(program);

    if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
        console.error('Unable to initialize the shader program: ' + gl.getProgramInfoLog(program));
        return null;
    }

    // Cleanup
    gl.deleteShader(computeShader);

    return program;
}

// Create and bind buffers
function createBuffers(gl, size, dataA, dataB) {
  // Input A
  const bufferA = gl.createBuffer();
  gl.bindBuffer(gl.SHADER_STORAGE_BUFFER, bufferA);
  gl.bufferData(gl.SHADER_STORAGE_BUFFER, dataA, gl.STATIC_DRAW);

  // Input B
  const bufferB = gl.createBuffer();
  gl.bindBuffer(gl.SHADER_STORAGE_BUFFER, bufferB);
  gl.bufferData(gl.SHADER_STORAGE_BUFFER, dataB, gl.STATIC_DRAW);

  // Output C
  const bufferC = gl.createBuffer();
  gl.bindBuffer(gl.SHADER_STORAGE_BUFFER, bufferC);
  gl.bufferData(gl.SHADER_STORAGE_BUFFER, size * 4, gl.STATIC_DRAW);
  // Note:  size * 4 because we are using floats, each of which are 4 bytes
  return { bufferA, bufferB, bufferC };
}

// Set up storage buffer binding points
function bindBuffers(gl, program, bufferA, bufferB, bufferC) {
  gl.useProgram(program);
  // Bind buffers to the program
  gl.bindBufferBase(gl.SHADER_STORAGE_BUFFER, 0, bufferA);
  gl.bindBufferBase(gl.SHADER_STORAGE_BUFFER, 1, bufferB);
  gl.bindBufferBase(gl.SHADER_STORAGE_BUFFER, 2, bufferC);
}

// Run the compute shader
function runComputeShader(gl, program, numElements) {
  gl.useProgram(program);

  // Determine number of workgroups
  const workgroupSize = 64;
  const numWorkgroups = Math.ceil(numElements / workgroupSize);

  // Dispatch compute shader
  gl.dispatchCompute(numWorkgroups, 1, 1);

  // Ensure the compute shader has finished running
  gl.memoryBarrier(gl.SHADER_STORAGE_BARRIER_BIT);
}

// Get results
function getResults(gl, bufferC, numElements) {
  const results = new Float32Array(numElements);
  gl.bindBuffer(gl.SHADER_STORAGE_BUFFER, bufferC);
  gl.getBufferSubData(gl.SHADER_STORAGE_BUFFER, 0, results);
  return results;
}

// Main execution
function main() {
  const numElements = 1024;
  const dataA = new Float32Array(numElements);
  const dataB = new Float32Array(numElements);

  // Initialize input data
  for (let i = 0; i < numElements; i++) {
    dataA[i] = i;
    dataB[i] = 2 * i;
  }

  const program = createComputeProgram(gl, shaderSource);
  if (!program) {
      return;
  }

  const { bufferA, bufferB, bufferC } = createBuffers(gl, numElements * 4, dataA, dataB);

  bindBuffers(gl, program, bufferA, bufferB, bufferC);

  runComputeShader(gl, program, numElements);

  const results = getResults(gl, bufferC, numElements);

  console.log('Results:', results);
  // Verify Results
  let allCorrect = true;
  for (let i = 0; i < numElements; ++i) {
      if (results[i] !== dataA[i] + dataB[i]) {
          console.error(`Error at index ${i}: Expected ${dataA[i] + dataB[i]}, got ${results[i]}`);
          allCorrect = false;
          break;
      }
  }
  if(allCorrect) {
      console.log('All results are correct.');
  }

  // Clean up buffers
  gl.deleteBuffer(bufferA);
  gl.deleteBuffer(bufferB);
  gl.deleteBuffer(bufferC);
  gl.deleteProgram(program);
}

main();

            

              
                Copy